{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "bbff7856",
   "metadata": {},
   "source": [
    "# Document Intelligence with Azure Cognitive Search Skillset\n",
    "\n",
    "AI enrichment is the application of machine learning models over content that isn't full text searchable in its raw form. Through enrichment, analysis and inference are used to create searchable content and structure where none previously existed.\n",
    "In this notebook, we use an [AI enrichment pipeline with Azure Cognitive Search](https://learn.microsoft.com/en-us/azure/search/cognitive-search-concept-intro) to extract information and create new content from various formats of document files.\n",
    "After we extract the enriched documents, we define a Feathr feature, materialize it, and utilize it from NLP (Natural Language Processing) scenarios such as [question-answering](https://huggingface.co/docs/transformers/tasks/question_answering) and [summarization](https://huggingface.co/tasks/summarization).\n",
    "\n",
    "The overall workflow is:\n",
    "1. Deploy Azure Cognitive Search and Feathr resources.\n",
    "2. Prepare mixed media sample documents.\n",
    "3. Extract texts by using Azure Cognitive Search Skillset and store the results.\n",
    "4. Define a Feathr feature with the enriched documents, register the feature, and materialize it.\n",
    "5. Use the feature for NLP scenarios."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "7da244ca",
   "metadata": {},
   "source": [
    "## 1. Deployment\n",
    "\n",
    "### Deploy Azure Cognitive Search service\n",
    "Please follow [this link](https://learn.microsoft.com/en-us/azure/search/search-create-service-portal) to create the search service from the Azure portal.\n",
    "\n",
    "### Deploy Feathr resources\n",
    "Please follow [this link](https://feathr-ai.github.io/feathr/how-to-guides/azure-deployment-arm.html) to deploy necessary resources via ARM template and grant permissions to access them."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "d2b26860",
   "metadata": {},
   "source": [
    "## 2. Prepare Dataset\n",
    "In this notebook, we use mixed media documents from an [Azure sample repository](https://github.com/Azure-Samples/azure-search-knowledge-mining).\n",
    "\n",
    "1. Download document files from the *[https://github.com/Azure-Samples/azure-search-knowledge-mining/tree/main/sample_documents](https://github.com/Azure-Samples/azure-search-knowledge-mining/tree/main/sample_documents)*\n",
    "2. Create a container called `cogsearch` at the Storage Account under the Cognitive Search resource group you deployed\n",
    "3. Upload the document files to the container "
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "2817d55c",
   "metadata": {},
   "source": [
    "## 3. AI Enrichment by using Azure Cognitive Search Skillset\n",
    "\n",
    "Before move on to the following steps, let's go over some of the important concepts in [Azure Cognitive Search](https://learn.microsoft.com/en-us/azure/search/search-what-is-azure-search):\n",
    "* **Azure Cognitive Search** is a cloud search service that gives developers infrastructure, APIs, and tools for building a rich search experience over private, heterogeneous content in web, mobile, and enterprise applications.\n",
    "* **Indexing** is an intake process that loads content into your search service and makes it searchable. AI enrichment through cognitive skills is an extension of indexing. If your content needs image or language analysis before it can be indexed, AI enrichment can extract text embedded in application files, translate text, and also infer text and structure from non-text files by analyzing the content.\n",
    "* **Skillset** is a reusable resource in Azure Cognitive Search that's attached to an indexer. It contains one or more skills that call built-in AI or external custom processing over documents retrieved from an external data source.\n",
    "* **Knowledge store** is a data sink created by a Cognitive Search enrichment pipeline that stores AI-enriched content in tables and blob containers in Azure Storage for independent analysis or downstream processing in non-search scenarios like knowledge mining.\n",
    "\n",
    "In this notebook, we use two built-in skills:\n",
    "* Translation and language detection\n",
    "* Optical Character Recognition (OCR) that recognizes printed and handwritten text in binary files\n",
    "\n",
    "For more details about the built-in skillset, see [here](https://learn.microsoft.com/en-us/azure/search/cognitive-search-predefined-skills).\n",
    "\n",
    "![Cognitive Search](../images/cognitive-search-enrichment-architecture.png)\n",
    "\n",
    "\n",
    "### 3.1 Set parameters for connecting to the resources"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0beae6e1",
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "from pprint import pprint\n",
    "import requests\n",
    "import time"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "21f08bb0",
   "metadata": {},
   "outputs": [],
   "source": [
    "# TODO fill the values:\n",
    "SERVICE_NAME = None  # Search Service name and API key\n",
    "API_KEY = None\n",
    "STORAGE_NAME = None  # Storage account for Azure Cognitive Search datasource and knowledge store\n",
    "STORAGE_KEY = None\n",
    "\n",
    "CONTAINER_NAME = \"cogsearch\"\n",
    "API_VERSION = \"2021-04-30-Preview\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fb1d7729",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Storage account connection string\n",
    "storage_connection_str = f\"DefaultEndpointsProtocol=https;AccountName={STORAGE_NAME};AccountKey={STORAGE_KEY};EndpointSuffix=core.windows.net\""
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "ac4aa5d1",
   "metadata": {},
   "source": [
    "To verify the access to the storage account with `STORAGE_NAME` and `STORAGE_KEY` we set from the previous cell, let's list the document blob names in the `cogsearch` container. Please make sure you created the container and uploaded the documents from the section **2. Prepare Datasets**."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c16c250b",
   "metadata": {},
   "outputs": [],
   "source": [
    "from azure.storage.blob import ContainerClient\n",
    "\n",
    "container_client = ContainerClient.from_connection_string(\n",
    "    storage_connection_str,\n",
    "    container_name=CONTAINER_NAME,\n",
    ")\n",
    "\n",
    "# Get name of the blobs\n",
    "output = [blob.name for blob in container_client.list_blobs()]\n",
    "output"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "782f89cd",
   "metadata": {},
   "source": [
    "### 3.2 Create DataSource, Index, and Indexer\n",
    "\n",
    "In Azure Cognitive Search, AI enrichment processing occurs during indexing (or data ingestion). The pipeline consists of:\n",
    "* Data source\n",
    "* Skill set\n",
    "* Index, and\n",
    "* Indexer.\n",
    "\n",
    "In this notebook, we use Search REST APIs to create them. First, let's define some helper functions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4785b56a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Cognitive search endpoint\n",
    "endpoint = f\"https://{SERVICE_NAME}.search.windows.net\"\n",
    "\n",
    "# Cognitive search REST API\n",
    "headers = {\n",
    "    \"Content-Type\": \"application/json\",\n",
    "    \"api-key\": API_KEY,\n",
    "}\n",
    "\n",
    "# Cognitive search resource names\n",
    "datasource_name = f\"{CONTAINER_NAME}-ds\"\n",
    "skillset_name = f\"{CONTAINER_NAME}-ss\"\n",
    "index_name = f\"{CONTAINER_NAME}-idx\"\n",
    "indexer_name = f\"{CONTAINER_NAME}-idxr\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "72be1e7c",
   "metadata": {},
   "outputs": [],
   "source": [
    "def construct_url(\n",
    "    endpoint: str,\n",
    "    resource_type: str,\n",
    "    resource_name: str = None,\n",
    "    action: str = None,\n",
    "    api_version: str = API_VERSION,\n",
    "):\n",
    "    \"\"\"Construct url for REST API calls.\"\"\"\n",
    "    components = [endpoint, resource_type]\n",
    "    \n",
    "    if resource_name:\n",
    "        components.append(resource_name)\n",
    "    if action:\n",
    "        components.append(action)\n",
    "    \n",
    "    return \"/\".join(components) + f\"?api-version={api_version}\"\n",
    "\n",
    "\n",
    "def create_or_update_resource(resource_type: str, resource_name: str, resource_def: dict):\n",
    "    \"\"\"Create or update Azure Cognitive Search resources.\"\"\"\n",
    "    r = requests.put(\n",
    "        construct_url(endpoint, resource_type, resource_name, None, API_VERSION),\n",
    "        data=json.dumps(resource_def),\n",
    "        headers=headers,\n",
    "    )\n",
    "\n",
    "    # The request should return a status code of 201 confirming success.\n",
    "    if r.status_code == 201:\n",
    "        print(f\"Successfully created {resource_type} {resource_name}\")\n",
    "    elif r.status_code == 204:\n",
    "        print(f\"Successfully updated {resource_type} {resource_name}\")\n",
    "    else:\n",
    "        print(r.json()[\"error\"][\"message\"])\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d2aebc42",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Clean-up previously created cognitive search resources if already exists\n",
    "for resource, resource_name in {\n",
    "    \"datasources\": datasource_name,\n",
    "    \"skillsets\": skillset_name,\n",
    "    \"indexes\": index_name,\n",
    "    \"indexers\": indexer_name,\n",
    "}.items():\n",
    "    r = requests.delete(\n",
    "        construct_url(endpoint, resource, resource_name),\n",
    "        headers=headers,\n",
    "    )\n",
    "    if r.status_code == 204:\n",
    "        print(f\"{resource} {resource_name} is successfully deleted.\")\n",
    "    else:\n",
    "        print(r.json()[\"error\"][\"message\"])"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "5d5644ef",
   "metadata": {},
   "source": [
    "#### Create a DataSource"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dabe6336",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create a data source\n",
    "datasource_def = {\n",
    "    \"name\": datasource_name,\n",
    "    \"description\": \"Feathr with Cognitive Search example\",\n",
    "    \"type\": \"azureblob\",\n",
    "    \"credentials\": {\n",
    "        \"connectionString\": storage_connection_str,\n",
    "    },\n",
    "    \"container\": {\n",
    "        \"name\": CONTAINER_NAME,\n",
    "    },\n",
    "}\n",
    "\n",
    "create_or_update_resource(\"datasources\", datasource_name, datasource_def)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "1b4032b6",
   "metadata": {},
   "source": [
    "#### Create a Skillset\n",
    "\n",
    "Each skill executes on the content of the document. During processing, Azure Cognitive Search cracks each document to read content from different file formats. Found text originating in the source file is placed into a generated `content` field, one for each document.\n",
    "\n",
    "For more details, see [Cognitive Search predefined skills](\n",
    "https://learn.microsoft.com/en-us/azure/search/cognitive-search-predefined-skills)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b5c551d2",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create a skillset\n",
    "skillset_def = {\n",
    "    \"name\": skillset_name,\n",
    "    \"description\": \"Apply OCR, \",\n",
    "    \"skills\": [\n",
    "        {\n",
    "            # Recognizes text and numbers in image files.\n",
    "            \"@odata.type\": \"#Microsoft.Skills.Vision.OcrSkill\",\n",
    "            \"context\": \"/document/normalized_images/*\",\n",
    "            \"defaultLanguageCode\": \"en\",\n",
    "            \"detectOrientation\": True,\n",
    "            \"inputs\": [\n",
    "                {\n",
    "                    \"name\": \"image\",\n",
    "                    \"source\": \"/document/normalized_images/*\"\n",
    "                }\n",
    "            ],\n",
    "            \"outputs\": [\n",
    "                {\n",
    "                    \"name\": \"text\"\n",
    "                }\n",
    "            ]\n",
    "        },\n",
    "        {\n",
    "            # Images and text are separated during the document cracking phase. The merge skill recombines them.\n",
    "            \"@odata.type\": \"#Microsoft.Skills.Text.MergeSkill\",\n",
    "            \"context\": \"/document\",\n",
    "            \"insertPreTag\": \" \",\n",
    "            \"insertPostTag\": \" \",\n",
    "            \"inputs\": [\n",
    "                {\n",
    "                    \"name\":\"text\", \n",
    "                    \"source\": \"/document/content\"\n",
    "                },\n",
    "                {\n",
    "                    \"name\": \"itemsToInsert\", \n",
    "                    \"source\": \"/document/normalized_images/*/text\"\n",
    "                },\n",
    "                {\n",
    "                    \"name\":\"offsets\", \n",
    "                    \"source\": \"/document/normalized_images/*/contentOffset\" \n",
    "                }\n",
    "            ],\n",
    "            \"outputs\": [\n",
    "                {\n",
    "                    \"name\": \"mergedText\", \n",
    "                    \"targetName\" : \"merged_text\"\n",
    "                }\n",
    "            ]\n",
    "        },\n",
    "        {\n",
    "            # Translates different language text to English.\n",
    "            \"@odata.type\": \"#Microsoft.Skills.Text.TranslationSkill\",\n",
    "            \"defaultToLanguageCode\": \"en\",\n",
    "            \"suggestedFrom\": \"es\",\n",
    "            \"context\": \"/document\",\n",
    "            \"inputs\": [\n",
    "                {\n",
    "                    \"name\": \"text\",\n",
    "                    \"source\": \"/document/merged_text\"\n",
    "                }\n",
    "            ],\n",
    "            \"outputs\": [\n",
    "                {\n",
    "                    \"name\": \"translatedText\",\n",
    "                    \"targetName\": \"translated_text\"\n",
    "                },\n",
    "                {\n",
    "                    \"name\": \"translatedFromLanguageCode\",\n",
    "                    \"targetName\": \"translated_from_language_code\"\n",
    "                },\n",
    "                {\n",
    "                    \"name\": \"translatedToLanguageCode\",\n",
    "                    \"targetName\": \"translated_to_language_code\"\n",
    "                }\n",
    "            ]\n",
    "        },\n",
    "        # Shaper skill to determine the schema and contents of the projection to Knowledge store.\n",
    "        {\n",
    "            \"@odata.type\": \"#Microsoft.Skills.Util.ShaperSkill\",\n",
    "            \"context\": \"/document\",\n",
    "            \"inputs\": [\n",
    "                {\n",
    "                    # metadata_storage_name is the document file (blob) name\n",
    "                    \"name\": \"metadata_storage_name\",\n",
    "                    \"source\": \"/document/metadata_storage_name\"\n",
    "                },\n",
    "                {\n",
    "                    \"name\": \"text\",\n",
    "                    \"source\": \"/document/translated_text\"\n",
    "                }\n",
    "            ],\n",
    "            \"outputs\": [\n",
    "                {\n",
    "                    \"name\": \"output\",\n",
    "                    \"targetName\": \"shaper_output\"\n",
    "                },\n",
    "            ],\n",
    "        },\n",
    "    ],\n",
    "    # Free skillset execution quota is 20 documents. To process more, set Azure Cognitive Service here.\n",
    "    \"cognitiveServices\": None,\n",
    "    # A knowledge store makes enriched content available in Azure Storage for downstream apps and workloads.\n",
    "    \"knowledgeStore\": {\n",
    "        \"storageConnectionString\": storage_connection_str,\n",
    "        \"projections\": [\n",
    "            {\n",
    "                # Store enrichment results into blobs:\n",
    "                \"objects\": [{\"storageContainer\": f\"{CONTAINER_NAME}output\", \"source\": \"/document/shaper_output\"}],\n",
    "            },\n",
    "        ],\n",
    "    },\n",
    "}\n",
    "\n",
    "create_or_update_resource(\"skillsets\", skillset_name, skillset_def)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "12b79360",
   "metadata": {},
   "source": [
    "#### Create an Index\n",
    "\n",
    "Provide the schema of the search index. A fields collection requires one field to be designated as the key. For blob content, this field is often the `metadata_storage_path` that uniquely identifies each blob in the container.\n",
    "\n",
    "In this schema, the \"text\" field receives OCR output, \"raw_content\" receives merged output, and \"content\" receives translation output."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "be5585b0",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create an index\n",
    "index_def = {\n",
    "    \"name\": index_name,\n",
    "    \"fields\": [\n",
    "        {\n",
    "            \"name\": \"text\",\n",
    "            \"type\": \"Collection(Edm.String)\",\n",
    "            \"searchable\": True,\n",
    "            \"sortable\": False,\n",
    "            \"filterable\": True,\n",
    "            \"facetable\": False\n",
    "        },\n",
    "        {\n",
    "            \"name\": \"content\",\n",
    "            \"type\": \"Edm.String\",\n",
    "            \"searchable\": True,\n",
    "            \"sortable\": False,\n",
    "            \"filterable\": False,\n",
    "            \"facetable\": False\n",
    "        },\n",
    "        {\n",
    "            \"name\": \"raw_content\",\n",
    "            \"type\": \"Edm.String\",\n",
    "            \"searchable\": False,\n",
    "            \"sortable\": False,\n",
    "            \"filterable\": False,\n",
    "            \"facetable\": False\n",
    "        },\n",
    "        {\n",
    "            \"name\": \"metadata_storage_path\",\n",
    "            \"type\": \"Edm.String\",\n",
    "            \"key\": True,\n",
    "            \"searchable\": True,\n",
    "            \"sortable\": False,\n",
    "            \"filterable\": False,\n",
    "            \"facetable\": False\n",
    "        },\n",
    "        {\n",
    "            \"name\": \"metadata_storage_name\",\n",
    "            \"type\": \"Edm.String\",\n",
    "            \"searchable\": True,\n",
    "            \"sortable\": False,\n",
    "            \"filterable\": True,\n",
    "            \"facetable\": False\n",
    "        }\n",
    "    ]\n",
    "}\n",
    "\n",
    "create_or_update_resource(\"indexes\", index_name, index_def)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "e21dc9c7",
   "metadata": {},
   "source": [
    "#### Create and Run an Indexer\n",
    "\n",
    "Creating an indexer invokes the pipeline."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1dc507f3",
   "metadata": {},
   "outputs": [],
   "source": [
    "indexer_def = {\n",
    "    \"name\": indexer_name,\n",
    "    \"dataSourceName\": datasource_name,\n",
    "    \"targetIndexName\": index_name,\n",
    "    \"skillsetName\": skillset_name,\n",
    "    \"cache\": {\n",
    "        \"enableReprocessing\": True,\n",
    "        \"storageConnectionString\": storage_connection_str,\n",
    "    },\n",
    "    # fieldMappings are processed before the skillset, sending content from the data source to target fields in an index.\n",
    "    \"fieldMappings\": [\n",
    "        {\n",
    "            \"sourceFieldName\": \"metadata_storage_path\",\n",
    "            \"targetFieldName\": \"metadata_storage_path\",\n",
    "            \"mappingFunction\": {\"name\": \"base64Encode\"}\n",
    "        },\n",
    "        {\n",
    "            \"sourceFieldName\": \"metadata_storage_name\",\n",
    "            \"targetFieldName\": \"metadata_storage_name\"\n",
    "        }\n",
    "    ],\n",
    "    # outputFieldMappings are for fields created by skills, after skillset execution.\n",
    "    # The references to sourceFieldName in outputFieldMappings don't exist until document cracking or enrichment creates them.\n",
    "    # The targetFieldName is a field in an index, defined in the index schema.\n",
    "    \"outputFieldMappings\": [\n",
    "        {\n",
    "            \"sourceFieldName\": \"/document/merged_text\",\n",
    "            \"targetFieldName\": \"raw_content\"\n",
    "        },\n",
    "        {\n",
    "            \"sourceFieldName\": \"/document/translated_text\",\n",
    "            \"targetFieldName\": \"content\"\n",
    "        },\n",
    "        {\n",
    "            \"sourceFieldName\": \"/document/normalized_images/*/text\",\n",
    "            \"targetFieldName\": \"text\"\n",
    "        }\n",
    "    ],\n",
    "    \"parameters\":\n",
    "    {\n",
    "        \"batchSize\": 1,\n",
    "        \"maxFailedItems\": -1,  # -1 to ignore errors during data import\n",
    "        \"maxFailedItemsPerBatch\": -1,\n",
    "        \"configuration\": \n",
    "        {\n",
    "            \"dataToExtract\": \"contentAndMetadata\",  # automatically extract the content from different file formats as well as metadata related to each file\n",
    "            \"imageAction\": \"generateNormalizedImages\"  # combined with the OCR Skill and Text Merge Skill, tells the indexer to extract text from the images\n",
    "        }\n",
    "    }   \n",
    "}\n",
    "\n",
    "create_or_update_resource(\"indexers\", indexer_name, indexer_def)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7b6b1435",
   "metadata": {},
   "source": [
    "Previous cell will create the indexer and run it. Now, we wait until the indexer run and check the result after done."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "58b27a5e",
   "metadata": {},
   "outputs": [],
   "source": [
    "while True:\n",
    "    r = requests.get(\n",
    "        construct_url(endpoint, \"indexers\", indexer_name, action=\"status\"),\n",
    "        headers=headers,\n",
    "    )\n",
    "    \n",
    "    try:\n",
    "        if r.json()[\"lastResult\"][\"endTime\"] is not None:\n",
    "            break\n",
    "    except:\n",
    "        pass\n",
    "\n",
    "    time.sleep(3)\n",
    "\n",
    "pprint(json.dumps(r.json(), indent=1))"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "8b30ead7",
   "metadata": {},
   "source": [
    "#### Test Search\n",
    "\n",
    "With REST API call, we can verify the result of cognitive search skills. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "14c46d62",
   "metadata": {},
   "outputs": [],
   "source": [
    "r = requests.post(\n",
    "    construct_url(\n",
    "        endpoint=endpoint,\n",
    "        resource_type=\"indexes\",\n",
    "        resource_name=index_name + \"/docs\",\n",
    "        action=\"search\",\n",
    "    ),\n",
    "    data=json.dumps({\n",
    "        \"search\": \"*\",\n",
    "        \"filter\": \"metadata_storage_name eq 'Cognitive Services and Bots  (spanish).pdf'\",  #'Mesh_for_Microsoft_Teams.docx'\",\n",
    "        \"select\": \"raw_content, content\",\n",
    "    }),\n",
    "    headers=headers,\n",
    ")\n",
    "\n",
    "result = r.json()['value'][0]\n",
    "print(\n",
    "    f\"[Content]{result['content'][:100]}\",\n",
    "    f\"[Raw content]{result['raw_content'][:100]}\",\n",
    "    sep=\"\\n=========\\n\",\n",
    ")\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "677fba41",
   "metadata": {},
   "source": [
    "## 4. Feathr Feature Store"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3879532d",
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "import os\n",
    "from pathlib import Path\n",
    "import shutil\n",
    "\n",
    "from pyspark.sql import DataFrame\n",
    "\n",
    "import feathr\n",
    "from feathr import (\n",
    "    FeathrClient,\n",
    "    # Feature data types\n",
    "    STRING, ValueType,\n",
    "    # Feature data sources\n",
    "    HdfsSource,\n",
    "    # Feature key\n",
    "    TypedKey,\n",
    "    # Feature types and anchor\n",
    "    Feature, FeatureAnchor,\n",
    "    # Materialization\n",
    "    MaterializationSettings, RedisSink,\n",
    "    # Offline feature computation\n",
    "    FeatureQuery, ObservationSettings,\n",
    ")\n",
    "from feathr.spark_provider.feathr_configurations import SparkExecutionConfiguration\n",
    "from feathr.utils.config import generate_config\n",
    "from feathr.utils.job_utils import get_result_df\n",
    "from feathr.utils.platform import is_databricks\n",
    "\n",
    "print(f\"Feathr version: {feathr.__version__}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "85d60f64",
   "metadata": {},
   "outputs": [],
   "source": [
    "RESOURCE_PREFIX = None  # TODO fill the value used to deploy the Feathr resources via ARM template\n",
    "PROJECT_NAME = \"cogsearch\"\n",
    "\n",
    "# Currently support: 'azure_synapse', 'databricks', and 'local' \n",
    "SPARK_CLUSTER = \"local\"\n",
    "\n",
    "# TODO fill values to use databricks cluster:\n",
    "DATABRICKS_CLUSTER_ID = None             # Set Databricks cluster id to use an existing cluster\n",
    "if is_databricks():\n",
    "    # If this notebook is running on Databricks, its context can be used to retrieve token and instance URL\n",
    "    ctx = dbutils.notebook.entry_point.getDbutils().notebook().getContext()\n",
    "    DATABRICKS_WORKSPACE_TOKEN_VALUE = ctx.apiToken().get()\n",
    "    SPARK_CONFIG__DATABRICKS__WORKSPACE_INSTANCE_URL = f\"https://{ctx.tags().get('browserHostName').get()}\"\n",
    "else:\n",
    "    DATABRICKS_WORKSPACE_TOKEN_VALUE = None                  # Set Databricks workspace token to use databricks\n",
    "    SPARK_CONFIG__DATABRICKS__WORKSPACE_INSTANCE_URL = None  # Set Databricks workspace url to use databricks\n",
    "\n",
    "# TODO fill values to use Azure Synapse cluster:\n",
    "AZURE_SYNAPSE_SPARK_POOL = None  # Set Azure Synapse Spark pool name\n",
    "AZURE_SYNAPSE_URL = None         # Set Azure Synapse workspace url to use Azure Synapse\n",
    "ADLS_KEY = None                  # Set Azure Data Lake Storage key to use Azure Synapse\n",
    "\n",
    "# If set True, use an interactive browser authentication to get the redis password.\n",
    "USE_CLI_AUTH = False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2765a239",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Use dbfs if the notebook is running on Databricks\n",
    "if is_databricks():\n",
    "    WORKING_DIR = f\"/dbfs/{PROJECT_NAME}\"\n",
    "else:\n",
    "    WORKING_DIR = PROJECT_NAME"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f10d4f66",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get an authentication credential to access Azure resources and register features\n",
    "if USE_CLI_AUTH:\n",
    "    # Use AZ CLI interactive browser authentication\n",
    "    !az login --use-device-code\n",
    "    from azure.identity import AzureCliCredential\n",
    "    credential = AzureCliCredential(additionally_allowed_tenants=['*'],)\n",
    "elif \"AZURE_TENANT_ID\" in os.environ and \"AZURE_CLIENT_ID\" in os.environ and \"AZURE_CLIENT_SECRET\" in os.environ:\n",
    "    # Use Environment variable secret\n",
    "    from azure.identity import EnvironmentCredential\n",
    "    credential = EnvironmentCredential()\n",
    "else:\n",
    "    # Try to use the default credential\n",
    "    from azure.identity import DefaultAzureCredential\n",
    "    credential = DefaultAzureCredential(\n",
    "        exclude_interactive_browser_credential=False,\n",
    "        additionally_allowed_tenants=['*'],\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ac848162",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Redis password\n",
    "if \"REDIS_PASSWORD\" not in os.environ:\n",
    "    from azure.keyvault.secrets import SecretClient\n",
    "    vault_url = f\"https://{RESOURCE_PREFIX}kv.vault.azure.net\"\n",
    "    secret_client = SecretClient(vault_url=vault_url, credential=credential)\n",
    "    retrieved_secret = secret_client.get_secret('FEATHR-ONLINE-STORE-CONN').value\n",
    "    os.environ['REDIS_PASSWORD'] = retrieved_secret.split(\",\")[1].split(\"password=\", 1)[1]"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "e1a4a74d",
   "metadata": {},
   "source": [
    "### 4.1 Initialize Feathr client"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "18c3820d",
   "metadata": {},
   "outputs": [],
   "source": [
    "config_path = generate_config(\n",
    "    resource_prefix=RESOURCE_PREFIX,\n",
    "    project_name=PROJECT_NAME,\n",
    "    spark_config__spark_cluster=SPARK_CLUSTER,\n",
    "    spark_config__azure_synapse__dev_url=AZURE_SYNAPSE_URL,\n",
    "    spark_config__azure_synapse__pool_name=AZURE_SYNAPSE_SPARK_POOL,\n",
    "    spark_config__databricks__workspace_instance_url=SPARK_CONFIG__DATABRICKS__WORKSPACE_INSTANCE_URL,\n",
    "    databricks_cluster_id=DATABRICKS_CLUSTER_ID,\n",
    ")\n",
    "\n",
    "with open(config_path, 'r') as f: \n",
    "    print(f.read())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "517c3357",
   "metadata": {},
   "outputs": [],
   "source": [
    "client = FeathrClient(config_path=config_path, credential=credential)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "b13f1329",
   "metadata": {},
   "source": [
    "### 4.2 Prepare Dataset\n",
    "\n",
    "To generate the data source for the features, let's get the Cognitive Search's AI skillset outputs from the Knowledge Store.\n",
    "As we defined the Knowledge Store projection to be objects, the outputs are stored as Json records. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ac2fd94a",
   "metadata": {},
   "outputs": [],
   "source": [
    "if \"spark\" in locals() or \"spark\" in globals():\n",
    "    spark.conf.set(f\"fs.azure.account.key.{STORAGE_NAME}.blob.core.windows.net\", STORAGE_KEY)\n",
    "else:\n",
    "    from pyspark.sql import SparkSession\n",
    "    spark = (\n",
    "        SparkSession\n",
    "        .builder\n",
    "        .appName(\"feathr\")\n",
    "        .config(\n",
    "            \"spark.jars.packages\",\n",
    "            \",\".join([\n",
    "                \"org.apache.spark:spark-avro_2.12:3.3.0\",\n",
    "                \"io.delta:delta-core_2.12:2.1.1\",\n",
    "                \"org.apache.hadoop:hadoop-azure:3.3.0\",\n",
    "                \"com.microsoft.azure:azure-storage:8.6.6\",\n",
    "            ])\n",
    "        )\n",
    "        .config(f\"fs.azure.account.key.{STORAGE_NAME}.blob.core.windows.net\", STORAGE_KEY)\n",
    "        .config(\"spark.sql.extensions\", \"io.delta.sql.DeltaSparkSessionExtension\")\n",
    "        .config(\"spark.sql.catalog.spark_catalog\", \"org.apache.spark.sql.delta.catalog.DeltaCatalog\")\n",
    "        .config(\"spark.ui.port\", \"8080\")  # Set ui port other than the default one (4040) so that feathr spark job doesn't fail. \n",
    "        .getOrCreate()\n",
    "    )\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "264b7a0a",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "# Read all the json records of the AI enrichment output\n",
    "df = spark.read.option(\"recursiveFileLookup\", \"true\").json(f\"wasbs://{CONTAINER_NAME}output@{STORAGE_NAME}.blob.core.windows.net/\")\n",
    "df.limit(5).toPandas()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a00d6dd6",
   "metadata": {},
   "outputs": [],
   "source": [
    "data_file_path = f\"{WORKING_DIR}/documents.parquet\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9069a239",
   "metadata": {},
   "outputs": [],
   "source": [
    "if Path(data_file_path).exists():\n",
    "    print(f\"Remove existing data file: {data_file_path}\")\n",
    "    shutil.rmtree(data_file_path)\n",
    "print(f\"Write data file to: {data_file_path}\")\n",
    "df.write.parquet(data_file_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "02923a12",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Upload files to cloud if needed\n",
    "if client.spark_runtime == \"local\":\n",
    "    # In local mode, we can use the same data path as the source.\n",
    "    data_source_path = data_file_path\n",
    "elif client.spark_runtime == \"databricks\" and is_databricks():\n",
    "    # If the notebook is running on databricks, we can use the same data path as the source.\n",
    "    data_source_path = data_file_path.replace(\"/dbfs\", \"dbfs:\")\n",
    "else:\n",
    "    # Otherwise, upload the local file to the cloud storage (either dbfs or adls).\n",
    "    data_source_path = client.feathr_spark_launcher.upload_or_get_cloud_path(data_file_path)    "
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "61c0df09",
   "metadata": {},
   "source": [
    "### 4.3 Define Features with UDF (User Defined Function)\n",
    "\n",
    "We preprocess the texts so that the later NLP models can consume them."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5794eda7",
   "metadata": {},
   "outputs": [],
   "source": [
    "def preprocessing(df: DataFrame) -> DataFrame:\n",
    "    import pyspark.sql.functions as F\n",
    "    \n",
    "    # Any types of text preprocessing\n",
    "    return df.withColumn(\"text\", F.regexp_replace(\"text\", \"\\n\", \" \"))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "54efd59e",
   "metadata": {},
   "source": [
    "Now, define features using the UDF."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5c4bdb6b",
   "metadata": {},
   "outputs": [],
   "source": [
    "hdfs_source = HdfsSource(\n",
    "    name=\"ai_enrichment\",\n",
    "    path=data_source_path,\n",
    "    preprocessing=preprocessing,\n",
    ")\n",
    "\n",
    "# key is required for the features from non-INPUT_CONTEXT source\n",
    "key = TypedKey(\n",
    "    key_column=\"metadata_storage_name\",\n",
    "    key_column_type=ValueType.STRING,\n",
    "    description=\"Document name\",\n",
    "    full_name=f\"{PROJECT_NAME}.doc_name\",\n",
    ")\n",
    "\n",
    "features = [\n",
    "    Feature(\n",
    "        name=\"f_text\",\n",
    "        key=key,\n",
    "        feature_type=STRING,\n",
    "        transform=\"text\",\n",
    "    ),\n",
    "]\n",
    "\n",
    "feature_anchor = FeatureAnchor(\n",
    "    name=\"data_feature_anchor\",\n",
    "    source=hdfs_source,\n",
    "    features=features,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5079d649",
   "metadata": {},
   "outputs": [],
   "source": [
    "client.build_features(\n",
    "    anchor_list=[feature_anchor],\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "38393c24",
   "metadata": {},
   "outputs": [],
   "source": [
    "query = FeatureQuery(\n",
    "    feature_list=[\"f_text\"],\n",
    "    key=key,\n",
    ")\n",
    "\n",
    "settings = ObservationSettings(\n",
    "    observation_path=data_source_path,\n",
    ")\n",
    "\n",
    "client.get_offline_features(\n",
    "    observation_settings=settings,\n",
    "    feature_query=query,\n",
    "    # For more details, see https://feathr-ai.github.io/feathr/how-to-guides/feathr-job-configuration.html\n",
    "    execution_configurations=SparkExecutionConfiguration({\n",
    "        \"spark.feathr.outputFormat\": \"parquet\",\n",
    "    }),\n",
    "    output_path=\"./text_features.parquet\",\n",
    ")\n",
    "\n",
    "client.wait_job_to_finish(timeout_sec=5000)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8692f2df",
   "metadata": {},
   "outputs": [],
   "source": [
    "get_result_df(client, data_format=\"parquet\").head(5)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "49ea7f5b",
   "metadata": {},
   "source": [
    "### 4.4 Register Features"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "try:\n",
    "    client.register_features()\n",
    "except Exception as e:\n",
    "    print(e)  \n",
    "print(client.list_registered_features(project_name=PROJECT_NAME))\n",
    "# You can get the actual features too by calling client.get_features_from_registry(PROJECT_NAME)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "1007a547",
   "metadata": {},
   "source": [
    "### 4.5 Materialize Features to REDIS"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0a3a9fd8",
   "metadata": {},
   "outputs": [],
   "source": [
    "FEATURE_TABLE_NAME = \"text_features\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "90d0092d",
   "metadata": {},
   "outputs": [],
   "source": [
    "redis_sink = RedisSink(table_name=FEATURE_TABLE_NAME)\n",
    "\n",
    "settings = MaterializationSettings(\n",
    "    name=FEATURE_TABLE_NAME + \".job\",  # job name\n",
    "    sinks=[redis_sink],\n",
    "    feature_names=[\"f_text\"],\n",
    ")\n",
    "\n",
    "client.materialize_features(\n",
    "    settings=settings,\n",
    "    execution_configurations={\"spark.feathr.outputFormat\": \"parquet\"},\n",
    "    allow_materialize_non_agg_feature=True,\n",
    ")\n",
    "\n",
    "client.wait_job_to_finish(timeout_sec=5000)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2aa13fb1",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Note, to get a single key, you may use client.get_online_features instead\n",
    "materialized_feature_values = client.get_online_features(\n",
    "    feature_table=FEATURE_TABLE_NAME,\n",
    "    key=\"NYSE_LNKD_2015.PDF\",\n",
    "    feature_names=[\"f_text\"],\n",
    ")\n",
    "materialized_feature_values[0][:1000]"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "a74fcbe7",
   "metadata": {},
   "source": [
    "## 5. NLP Scenarios\n",
    "\n",
    "We use [HuggingFace Transformer package](https://huggingface.co/docs/transformers/installation) to demonstrate simple NLP scenarios with the materialized features. Specifically, we use [summarization and question-answering pipelines](https://huggingface.co/docs/transformers/main_classes/pipelines) that come with pre-trained models for the simplicity's sake.\n",
    "\n",
    "Firstly, install *Transformer* and *pyTorch* packages:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "45821650",
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install -U transformers torch --extra-index-url https://download.pytorch.org/whl/cu116"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2215cae8",
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import pipeline"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "7263b727",
   "metadata": {},
   "source": [
    "Now, let's use question-answering pipeline of the HuggingFace package."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d38b589a",
   "metadata": {},
   "outputs": [],
   "source": [
    "qa_model = pipeline(\"question-answering\")\n",
    "qa_model(\n",
    "    question=\"what is the LinkedIn's financial goal\",\n",
    "    context=materialized_feature_values[0],\n",
    ")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "6f6c5ee3",
   "metadata": {},
   "source": [
    "With the summarization pipeline,"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3594c55f",
   "metadata": {},
   "outputs": [],
   "source": [
    "summarizer = pipeline(\"summarization\")\n",
    "\n",
    "# The pre-trained model only accepts 1024 tokens as input and thus we set truncation=True\n",
    "summarizer(materialized_feature_values[0], truncation=True)\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "26b9927c",
   "metadata": {},
   "source": [
    "## 6. Advanced Topics\n",
    "\n",
    "In this notebook, we have gone through how to utilize Azure Cognitive Search Skills to extract and translate texts from various formats of documents and use them with Feathr Feature Store for NLP scenarios.\n",
    "\n",
    "Here is a list of advanced topics we did not cover from this notebook:\n",
    "\n",
    "* [Deploy a model to AKS (Azure Kubernetes Service) via AzureML (Azure Machine Learning) SDK](https://learn.microsoft.com/en-us/azure/machine-learning/v1/how-to-deploy-azure-kubernetes-service?tabs=python).\n",
    "* [Use Python and AI to generate searchable content from Azure blobs](https://learn.microsoft.com/en-us/azure/search/cognitive-search-tutorial-blob-python)\n",
    "* [Enrich cognitive search index with custom classes](https://learn.microsoft.com/en-us/azure/cognitive-services/language-service/custom-text-classification/tutorials/cognitive-search?tabs=multi-classification%2CLanguage-studio)\n",
    "* [Build and deploy a form recognizer custom skill](https://learn.microsoft.com/en-us/training/modules/build-form-recognizer-custom-skill-for-azure-cognitive-search/4-exercise-build-deploy)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6e907bea",
   "metadata": {},
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.9 (main, Dec 15 2022, 18:18:30) [Clang 14.0.0 (clang-1400.0.29.202)]"
  },
  "vscode": {
   "interpreter": {
    "hash": "aee8b7b246df8f9039afb4144a1f6fd8d2ca17a180786b69acc140d282b71a49"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
